#include <filesystem>
#include <nlohmann/json.hpp>
#include "InputData.h"
#include "CVUtils.h"

namespace fs = std::filesystem;
using namespace torch::indexing;
using json = nlohmann::json;

namespace ns { SInputData inputDataFromNerfStudio(const std::string& vProjectRoot); }
namespace cm { SInputData inputDataFromColmap(const std::string& vProjectRoot); }
namespace osfm { SInputData inputDataFromOpenSfM(const std::string& vProjectRoot); }

//******************************************************************
//FUNCTION:
SInputData inputDataFromX(const std::string& vProjectRoot) {
    fs::path Root(vProjectRoot);

    if (fs::exists(Root / "transforms.json")) {
        return ns::inputDataFromNerfStudio(vProjectRoot);
    }
    else if (fs::exists(Root / "sparse") || fs::exists(Root / "cameras.bin")) {
        return cm::inputDataFromColmap(vProjectRoot);
    }
    else if (fs::exists(Root / "reconstruction.json")) {
        return osfm::inputDataFromOpenSfM(vProjectRoot);
    }
    else if (fs::exists(Root / "opensfm" / "reconstruction.json")) {
        return osfm::inputDataFromOpenSfM((Root / "opensfm").string());
    }
    else {
        throw std::runtime_error("Invalid project folder (must be either a colmap or nerfstudio project folder)");
    }
}

//******************************************************************
//FUNCTION:
torch::Tensor CCamera::getIntrinsicsMatrix() {
    return torch::tensor({ {m_Fx, 0.0f, m_Cx},
                          {0.0f, m_Fy, m_Cy},
                          {0.0f, 0.0f, 1.0f} }, torch::kFloat32);
}

//******************************************************************
//FUNCTION:
void CCamera::loadImage(float vDownscaleFactor) {
    // Populates image and K, then updates the camera parameters
    // Caution: this function has destructive behaviors
    // and should be called only once
    if (m_Image.numel()) std::runtime_error("loadImage already called");
    std::cout << "Loading " << m_FilePath << std::endl;

    float ScaleFactor = 1.0f / vDownscaleFactor;
    cv::Mat cImg = imreadRGB(m_FilePath);

    float RescaleF = 1.0f;
    // If camera intrinsics don't match the image dimensions 
    if (cImg.rows != m_Height || cImg.cols != m_Width) {
        RescaleF = static_cast<float>(cImg.rows) / static_cast<float>(m_Height);
    }
    m_Fx *= ScaleFactor * RescaleF;
    m_Fy *= ScaleFactor * RescaleF;
    m_Cx *= ScaleFactor * RescaleF;
    m_Cy *= ScaleFactor * RescaleF;

    if (vDownscaleFactor > 1.0f) {
        float f = 1.0f / vDownscaleFactor;
        cv::resize(cImg, cImg, cv::Size(), f, f, cv::INTER_AREA);
    }

    m_K = getIntrinsicsMatrix();
    cv::Rect Roi;

    if (__hasDistortionParameters()) {
        // Undistort
        std::vector<float> distCoeffs = __undistortionParameters();
        cv::Mat CK = floatNxNtensorToMat(m_K);
        cv::Mat NewK = cv::getOptimalNewCameraMatrix(CK, distCoeffs, cv::Size(cImg.cols, cImg.rows), 0, cv::Size(), &Roi);

        cv::Mat Undistorted = cv::Mat::zeros(cImg.rows, cImg.cols, cImg.type());
        cv::undistort(cImg, Undistorted, CK, distCoeffs, NewK);

        m_Image = imageToTensor(Undistorted);
        m_K = floatNxNMatToTensor(NewK);
    }
    else {
        Roi = cv::Rect(0, 0, cImg.cols, cImg.rows);
        m_Image = imageToTensor(cImg);
    }

    // Crop to ROI
    m_Image = m_Image.index({ Slice(Roi.y, Roi.y + Roi.height), Slice(Roi.x, Roi.x + Roi.width), Slice() });

    // Update parameters
    m_Height = m_Image.size(0);
    m_Width = m_Image.size(1);
    m_Fx = m_K[0][0].item<float>();
    m_Fy = m_K[1][1].item<float>();
    m_Cx = m_K[0][2].item<float>();
    m_Cy = m_K[1][2].item<float>();
}

//******************************************************************
//FUNCTION:
torch::Tensor CCamera::getImage(int vDownscaleFactor) {
    if (vDownscaleFactor <= 1) return m_Image;
    else {
        if (m_ImagePyramids.find(vDownscaleFactor) != m_ImagePyramids.end()) {
            return m_ImagePyramids[vDownscaleFactor];
        }

        // Rescale, store and return
        cv::Mat cImg = tensorToImage(m_Image);
        cv::resize(cImg, cImg, cv::Size(cImg.cols / vDownscaleFactor, cImg.rows / vDownscaleFactor), 0.0, 0.0, cv::INTER_AREA);
        torch::Tensor T = imageToTensor(cImg);
        m_ImagePyramids[vDownscaleFactor] = T;
        return T;
    }
}

//******************************************************************
//FUNCTION:
bool CCamera::__hasDistortionParameters() {
    return m_K1 != 0.0f || m_K2 != 0.0f || m_K3 != 0.0f || m_P1 != 0.0f || m_P2 != 0.0f;
}

//******************************************************************
//FUNCTION:
std::vector<float> CCamera::__undistortionParameters() {
    std::vector<float> Parameters = { m_K1, m_K2, m_P1, m_P2, m_K3, 0.0f, 0.0f, 0.0f };
    return Parameters;
}

//******************************************************************
//FUNCTION:
std::tuple<std::vector<CCamera>, CCamera*> SInputData::getCameras(bool vValidate, const std::string& vValImage) {
    if (!vValidate) return std::make_tuple(_Cameras, nullptr);
    else {
        size_t ValIdx = -1;
        int RandomSeed = 42;
        std::srand(RandomSeed);

        if (vValImage == "random") {
            ValIdx = std::rand() % _Cameras.size();
        }
        else {
            for (size_t i = 0; i < _Cameras.size(); i++) {
                if (fs::path(_Cameras[i].getFilePath()).filename().string() == vValImage) {
                    ValIdx = i;
                    break;
                }
            }
            if (ValIdx == -1) throw std::runtime_error(vValImage + " not in the list of cameras");
        }

        std::vector<CCamera> Cams;
        CCamera* ValCam = nullptr;

        for (size_t i = 0; i < _Cameras.size(); i++) {
            if (i != ValIdx) Cams.push_back(_Cameras[i]);
            else ValCam = &_Cameras[i];
        }

        return std::make_tuple(Cams, ValCam);
    }
}

//******************************************************************
//FUNCTION:
void SInputData::saveCameras(const std::string& vFilename, bool vKeepCrs) {
    json j = json::array();

    for (size_t i = 0; i < _Cameras.size(); i++) {
        CCamera& Cam = _Cameras[i];

        json Camera = json::object();
        Camera["id"] = i;
        Camera["img_name"] = fs::path(Cam.getFilePath()).filename().string();
        Camera["width"] = Cam.getWidth();
        Camera["height"] = Cam.getHeight();
        Camera["fx"] = Cam.getFx();
        Camera["fy"] = Cam.getFy();

        torch::Tensor R = Cam.getCamToWorld().index({Slice(None, 3), Slice(None, 3)});
        torch::Tensor T = Cam.getCamToWorld().index({ Slice(None, 3), Slice(3,4) }).squeeze();

        // Flip z and y
        R = torch::matmul(R, torch::diag(torch::tensor({ 1.0f, -1.0f, -1.0f })));

        if (vKeepCrs) T = (T / _Scale) + _Translation;

        std::vector<float> Position(3);
        std::vector<std::vector<float>> Rotation(3, std::vector<float>(3));
        for (int i = 0; i < 3; i++) {
            Position[i] = T[i].item<float>();
            for (int j = 0; j < 3; j++) {
                Rotation[i][j] = R[i][j].item<float>();
            }
        }

        Camera["position"] = Position;
        Camera["rotation"] = Rotation;
        j.push_back(Camera);
    }

    std::ofstream of(vFilename);
    of << j;
    of.close();

    std::cout << "Wrote " << vFilename << std::endl;
}